# 基于规则的睡眠分析程序
# 功能：实时接收原始EEG数据，并使用基于频带功率的规则进行睡眠状态判断
# 优势：轻量、无需训练、实时性强

import numpy as np
import time
import csv
import sys
import os
from datetime import datetime
import math

# 添加父目录到路径以导入NeuroPy3
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from NeuroPy3 import NeuroPy

class SleepAnalyzer:
    def __init__(self, port="COM5", sfreq=512, log_step_sec=1):
        self.port = port
        self.sfreq = sfreq
        self.log_step_size = sfreq * log_step_sec  # 每秒记录日志
        
        self.log_counter = 0
        self.buffer = []
        self.poor_signal_buffer = []
        
        self.neuropy = NeuroPy(port, 57600)
        self.log_file = "mne_selfcreate.csv"
        
        # 添加一个变量来存储前一个状态，用于上下文判断
        self.previous_stage = "未知"
        self.state_duration = {}  # 记录每个状态的持续时间
        self.min_rem_duration = 5  # REM睡眠最少持续5分钟
    
    def _init_log(self):
        """初始化CSV日志文件"""
        headers = ['时间戳', '数据长度(秒)', '睡眠阶段', 'Delta功率(dB)', 'Theta功率(dB)', 'Alpha功率(dB)', 'Beta功率(dB)', 'Gamma功率(dB)', '信号质量']
        with open(self.log_file, 'w', newline='', encoding='utf-8') as f:
            writer = csv.writer(f)
            writer.writerow(headers)
    
    def _data_callback(self, raw_value):
        """接收原始EEG数据的回调函数"""
        self.buffer.append(raw_value)
    
    def _poor_signal_callback(self, poor_signal_value):
        """信号质量回调函数"""
        self.poor_signal_buffer.append(poor_signal_value)

    def _analyze_sleep_status(self, data_row, previous_state):
        """改进的睡眠状态分析算法"""
        try:
            delta, theta, alpha, beta, gamma, quality = map(float, data_row)
        except (ValueError, IndexError):
            return "数据格式错误"
    
        # 信号质量检查
        if quality != 0.00:
            return "信号异常"
        
        # 定义合理的状态转换规则
        valid_transitions = {
            '清醒': ['放松', '清醒'],
            '放松': ['清醒', '入睡', '放松'],
            '入睡': ['放松', '浅睡眠', '清醒'],
            '浅睡眠': ['入睡', '深度睡眠', 'REM睡眠', '清醒'],
            '深度睡眠': ['浅睡眠', 'REM睡眠'],
            'REM睡眠': ['浅睡眠', '清醒'],
            '未知': ['清醒']
        }
        
        # 计算相对功率
        total_power = delta + theta + alpha + beta + gamma
        if total_power <= 0:
            return "数据异常"
        
        alpha_ratio = alpha / total_power
        theta_ratio = theta / total_power
        delta_ratio = delta / total_power
        
        # 1. 清醒状态：高频活动且Alpha波相对较强
        if (beta + gamma) > (delta + theta) and alpha_ratio > 0.2:
            candidate_state = "清醒"
        
        # 2. 放松状态：Alpha波占主导
        elif alpha_ratio > 0.3 and alpha > beta:
            candidate_state = "放松"
        
        # 3. 入睡状态：从放松过渡，Theta波开始增强
        elif theta_ratio > alpha_ratio and theta_ratio > 0.25:
            candidate_state = "入睡"
        
        # 4. 浅睡眠：Theta波占主导
        elif theta_ratio > 0.35 and theta > alpha:
            candidate_state = "浅睡眠"
        
        # 5. 深度睡眠：Delta波显著增强
        elif delta_ratio > 0.4 and delta > theta:
            candidate_state = "深度睡眠"
        
        # 6. REM睡眠：只能从深度睡眠或浅睡眠转换而来
        elif (previous_state in ['深度睡眠', '浅睡眠'] and 
              (beta + gamma) > delta and theta > delta * 0.8):
            candidate_state = "REM睡眠"
        
        else:
            candidate_state = "清醒"  # 默认状态
        
        # 检查状态转换是否合理
        if previous_state in valid_transitions:
            if candidate_state in valid_transitions[previous_state]:
                return candidate_state
            else:
                # 如果转换不合理，保持前一状态或回到清醒
                if previous_state == '清醒':
                    return '清醒'
                else:
                    return '清醒'  # 异常情况回到清醒状态
        
        return candidate_state

    def _process_data(self):
        """处理缓冲区数据并进行分析"""
        buffer_length = len(self.buffer)
        
        # --- 每秒记录日志和进行分析 ---
        if buffer_length >= self.log_step_size * (self.log_counter + 1):
            self.log_counter += 1
            
            # 从NeuroPy设备获取实时的频带功率数据
            # 原始功率值，需要转换为dB
            raw_delta = self.neuropy.delta
            raw_theta = self.neuropy.theta
            raw_alpha = (self.neuropy.lowAlpha + self.neuropy.highAlpha) / 2
            raw_beta = (self.neuropy.lowBeta + self.neuropy.highBeta) / 2
            raw_gamma = (self.neuropy.lowGamma + self.neuropy.midGamma) / 2
            
            # 使用设备的实时信号质量值
            signal_quality = self.neuropy.poorSignal
            
            # 将原始功率值转换为dB，并处理可能为0的功率值
            delta_power = 10 * math.log10(raw_delta) if raw_delta > 0 else -999.0
            theta_power = 10 * math.log10(raw_theta) if raw_theta > 0 else -999.0
            alpha_power = 10 * math.log10(raw_alpha) if raw_alpha > 0 else -999.0
            beta_power = 10 * math.log10(raw_beta) if raw_beta > 0 else -999.0
            gamma_power = 10 * math.log10(raw_gamma) if raw_gamma > 0 else -999.0
            
            # 组装数据用于分析
            data_row = [delta_power, theta_power, alpha_power, beta_power, gamma_power, signal_quality]

            # 进行状态分析
            current_stage = self._analyze_sleep_status(data_row, self.previous_stage)
            
            # 更新前一状态
            self.previous_stage = current_stage

            # 记录日志
            print(f"第{self.log_counter}秒 - 状态: {current_stage}")
            self._log_result(1.0, current_stage, delta_power, theta_power, alpha_power, beta_power, gamma_power, signal_quality)

    def _log_result(self, length, stage, delta, theta, alpha, beta, gamma, quality):
        """记录分析结果到日志"""
        timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        with open(self.log_file, 'a', newline='', encoding='utf-8') as f:
            writer = csv.writer(f)
            writer.writerow([timestamp, length, stage, f"{delta:.6f}", f"{theta:.6f}", f"{alpha:.6f}", f"{beta:.6f}", f"{gamma:.6f}", f"{quality:.2f}"])
    
    def start(self):
        """启动数据采集"""
        self.neuropy.setCallBack("rawValue", self._data_callback)
        self.neuropy.setCallBack("poorSignal", self._poor_signal_callback)
        self.neuropy.start()
        print(f"设备已连接到 {self.port}，开始采集数据...")
        
        # 等待缓冲区累积到至少一个步长的数据
        while len(self.buffer) < self.log_step_size:
            time.sleep(0.1)
        
        try:
            while True:
                self._process_data()
                time.sleep(0.1)  # 短暂休眠避免CPU占用过高
        except KeyboardInterrupt:
            print("\n停止采集...")
        finally:
            self.neuropy.stop()
            print("设备已断开")

if __name__ == "__main__":
    analyzer = SleepAnalyzer(log_step_sec=1)
    analyzer.start()

# 基于您的数据范围调整的阈值
CONFIG = {
    'relaxed_alpha_threshold': 40.0,  # Alpha > 40dB 且相对较强时为放松
    'drowsy_theta_threshold': 0.25,   # Theta占比 > 25% 时为入睡
    'light_sleep_theta_ratio': 0.3,   # Theta占比 > 30% 时为浅睡眠
    'deep_sleep_delta_ratio': 0.35,   # Delta占比 > 35% 时为深度睡眠
}


def _validate_rem_sleep(self, current_state, previous_state):
    """验证REM睡眠的合理性"""
    if current_state == "REM睡眠":
        # REM睡眠必须从深度睡眠转换而来，且有足够的睡眠时间
        if previous_state not in ['深度睡眠', '浅睡眠']:
            return False
        
        # 检查是否有足够的睡眠时间（简化版本）
        total_sleep_time = (self.state_duration.get('浅睡眠', 0) + 
                            self.state_duration.get('深度睡眠', 0))
        if total_sleep_time < 30:  # 至少30分钟睡眠后才可能进入REM
            return False
    
    return True


def _detect_eye_closure_transition(self, current_data, previous_data):
    """检测闭眼转换"""
    # 计算Alpha波的相对变化
    alpha_change = (current_data['alpha'] - previous_data['alpha']) / previous_data['alpha']
    
    # 闭眼时Alpha波通常会增强10-30%
    if alpha_change > 0.1 and current_data['alpha'] > current_data['beta']:
        return "闭眼放松"
    
    return None

def _analyze_relaxation_depth(self, data_row):
    """分析放松深度"""
    delta, theta, alpha, beta, gamma = data_row[:5]
    
    # 计算放松指数
    relaxation_index = (alpha + theta) / (beta + gamma)
    
    if relaxation_index > 1.5:
        return "深度放松"
    elif relaxation_index > 1.2:
        return "中度放松"
    elif relaxation_index > 1.0:
        return "轻度放松"
    else:
        return "清醒"
